#!/usr/bin/env python3
import cv2
import apriltag
import cv_bridge
import rclpy
import numpy as np
from rclpy.node import Node
from sensor_msgs.msg import Image

camera_matrix = np.array(
    [
        [1605.569046, 0.000000, 678.129552],
        [0.000000, 1602.841388, 372.096778],
        [0.000000, 0.000000, 1.000000],
    ]
)
dist_coeffs = np.array([0.102228, -0.111735, 0.001705, -0.000391])
cam_params = [1605.569046, 1602.841388, 678.129552, 372.096778]


class AprilTagFinder(Node):
    def __init__(self):
        super().__init__("april_tag_finder")
        self.publisher = self.create_publisher(Image, "image_processed", 1)
        self.subscription = self.create_subscription(
            Image, "camera_image", self.image_callback, 1
        )
        self.tag_detector = apriltag.Detector(
            apriltag.DetectorOptions(families="tag36h11")
        )

    def image_callback(self, msg):
        cv_image = cv_bridge.CvBridge().imgmsg_to_cv2(msg, desired_encoding="bgr8")
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
        tags = self.tag_detector.detect(gray)
        for tag in tags:
            # self.get_logger().info(f"id:{tag.tag_id}")
            for i in range(4):
                cv2.circle(
                    cv_image, tuple(tag.corners[i].astype(int)), 4, (255, 0, 0), 2
                )
            cv2.circle(cv_image, tuple(tag.center.astype(int)), 4, (2, 180, 200), 4)
            M, e1, e2 = self.tag_detector.detection_pose(tag, cam_params)
            P = M[:3, :4]  # 相机投影矩阵
            P = np.matmul(camera_matrix, P)
            x = np.matmul(P, np.array([[1], [0], [0], [1]]))
            x = x / x[2]
            y = np.matmul(P, np.array([[0], [1], [0], [1]]))
            y = y / y[2]
            z = np.matmul(P, np.array([[0], [0], [-1], [1]]))
            z = z / z[2]
            cv2.line(
                cv_image,
                tuple(tag.center.astype(int)),
                tuple(x[:2].reshape(-1).astype(int)),
                (0, 0, 255),
                2,
            )  # x轴红色
            cv2.line(
                cv_image,
                tuple(tag.center.astype(int)),
                tuple(y[:2].reshape(-1).astype(int)),
                (0, 255, 0),
                2,
            )  # y轴绿色
            cv2.line(
                cv_image,
                tuple(tag.center.astype(int)),
                tuple(z[:2].reshape(-1).astype(int)),
                (255, 0, 0),
                2,
            )  # z轴蓝色

        self.publisher.publish(cv_bridge.CvBridge().cv2_to_imgmsg(cv_image, "bgr8"))


def main(args=None):
    rclpy.init(args=args)
    node = AprilTagFinder()
    node.get_logger().info("AprilTag Finder node started")
    rclpy.spin(node)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
